import pandas as pd
import numpy as np


def mark_d_ma(df):
    ma_df = df.copy()
    
    ma_df = ma_df.sort_values(['code', 'date'])
    ma_df['date'] = pd.to_datetime(ma_df['date'])
    
    # 合并分组操作（关键优化点）[7](@ref)
    grouped = ma_df.groupby('code')['close']
    
    # 批量计算所有移动平均（含新增ma8/ma10）[2,4](@ref)
    ma_windows = [3, 5, 8, 10, 15,120,250]
    for window in ma_windows:
        ma_df[f'ma{window}'] = grouped.transform(lambda s: s.rolling(window, min_periods=1).mean())
    
    # 计算乖离率（向量化实现）[2,5](@ref)
    for window in [3, 5, 15,120,250]:
        ma = ma_df[f'ma{window}']
        ma_df[f'bias{window}'] = (ma_df['close'] - ma) / ma * 100
    
    # 计算倾斜率（差分法优化）[6,8](@ref)
    for window in [3, 5, 15,120,250]:
        ma_col = f'ma{window}'
        # 计算当日与前一日MA的差值作为短期斜率
        ma_df[f'slope{window}'] = ma_df.groupby('code')[ma_col].transform(lambda s: s.diff())
    
    # 计算多均线条件（向量化布尔运算）[3](@ref)
    ma_df['d_duo'] = (
        (ma_df['ma3'] > ma_df['ma5']) & 
        (ma_df['ma5'] >= ma_df['ma8']) & 
        (ma_df['ma8'] >= ma_df['ma10'])
    ).astype(int)  # 转为1/0便于后续计算
    ma_df['d_kong'] = (
        (ma_df['ma3'] < ma_df['ma5']) & 
        (ma_df['ma5'] < ma_df['ma8']) & 
        (ma_df['ma8'] < ma_df['ma10'])
    ).astype(int)  # 转为1/0便于后续计算
    # 删除中间列（可选）
    ma_df.drop(['ma3','ma5','ma8','ma10','ma15','ma120','ma250',], axis=1, inplace=True)
    ma_df = drop_base_columns(ma_df)
    return ma_df.round(2)


# def mark_w_ma(days_df, w_df):
#     """
#     安全计算动态周均线方法（修复排序问题+避免未来数据）
    
#     参数：
#         days_df : 日线数据（date, code, close）
#         w_df    : 周线数据（date列表示当周最后一个交易日，code, close）
    
#     返回：
#         含3/5/15周均线的DataFrame
#     """
#     # ========== 阶段1：数据预处理 ==========
#     # 处理日线数据
#     df = days_df.copy()
#     df['date'] = pd.to_datetime(df['date'])
    
#     # 处理周线数据（关键修复点）
#     w_df = (w_df.rename(columns={'date': 'week_end_date'})
#            .sort_values(['code', 'week_end_date'])  # 必须排序
#            .assign(week_end_date=lambda x: pd.to_datetime(x.week_end_date)))
    
#     # ========== 阶段2：安全合并周收盘价 ==========
#     df = pd.merge_asof(
#         df.sort_values('date'),  # 左表必须排序
#         w_df[['code', 'week_end_date', 'close']]
#             .rename(columns={'close': 'dynamic_week_close'}).sort_values('week_end_date'),
#         left_on='date',
#         right_on='week_end_date',
#         by='code',
#         direction='backward'  # 关键参数：只用历史周数据
#     ).drop(columns='week_end_date')
    
#     # ========== 阶段3：滚动计算周均线 ==========
#     # 按code分组后按日期排序
#     df = df.sort_values(['code', 'date'])
    
#     ma_configs = [
#         ('ma3w', 3),   # 3周均线
#         ('ma5w', 5),   # 5周均线
#         ('ma8w', 8),   # 8周均线
#         ('ma15w', 15)  # 15周均线
#     ]
    
#     for col, window in ma_configs:
#         df[col] = df.groupby('code', group_keys=False)['dynamic_week_close'].transform(
#             lambda s: s.rolling(
#                 window=window,
#                 min_periods=1,
#                 closed='left'  # 关键参数：排除当前周
#             ).mean()
#         )
        
#     # 阶段5：乖离率计算
#     for window in [3,5,15]:
#         ma_col = f'ma{window}w'
#         df[f'bias{window}w'] = (df['close'] - df[ma_col]) / df[ma_col] * 100
   
    
#     # 阶段6：倾斜率计算（按周数窗口）
#     df = calculate_slopes(df)

    
#     # 阶段7：均线多空判断（严格对齐需求）
#     df['w_duo1'] = (df['ma3w'] > df['ma5w']) & (df['ma5w'] >= df['ma8w'])
#     df['w_duo2'] = (df['ma3w'] > df['ma5w']) & (df['ma5w'] >= df['ma15w'])
#     df.drop(['dynamic_week_close'], axis=1, inplace=True)

#     df = drop_base_columns(df)

    
#     return df.round(2)

def mark_w_ma(days_df, w_df):
    """
    修复版动态周均线计算方法
    主要修复点：
    1. 处理本周尚未结束时的动态周收盘价计算
    2. 确保滚动计算不包含未来数据
    """
    # ========== 阶段1：数据预处理 ==========
    # 处理日线数据
    df = days_df.copy()
    df['date'] = pd.to_datetime(df['date'])
    
    # 处理周线数据
    w_df = w_df.copy()
    w_df['date'] = pd.to_datetime(w_df['date'])
    w_df = w_df.rename(columns={'date': 'week_end_date', 'close': 'week_close'})
    
    # 添加自然周标记
    df['year_week'] = df['date'].dt.isocalendar().year.astype(str) + '-' + df['date'].dt.isocalendar().week.astype(str).str.zfill(2)
    
    # ========== 阶段2：安全合并周收盘价 ==========
    # 合并周线收盘价（仅用于已结束的周）
    df = pd.merge_asof(
        df.sort_values('date'),
        w_df[['code', 'week_end_date', 'week_close']].sort_values('week_end_date'),
        left_on='date',
        right_on='week_end_date',
        by='code',
        direction='backward'
    )
    
    # ========== 阶段3：动态计算周收盘价 ==========
    # 核心修复：当本周尚未结束时，使用当日收盘价作为动态周收盘价
    df['dynamic_week_close'] = np.where(
        df['week_end_date'] >= df['date'],  # 本周尚未结束
        df['close'],  # 使用当日收盘价
        df['week_close']  # 使用历史周收盘价
    )
    
    # ========== 阶段4：滚动计算周均线 ==========
    # 按code分组后按日期排序
    df = df.sort_values(['code', 'date'])
    
    ma_configs = [
        ('ma3w', 3),   # 3周均线
        ('ma5w', 5),   # 5周均线
        ('ma8w', 8),   # 8周均线
        ('ma15w', 15)  # 15周均线
    ]
    
    for col, window in ma_configs:
        # 关键修复：使用shift确保不包含当前周
        df[col] = df.groupby('code', group_keys=False)['dynamic_week_close'].transform(
            lambda s: s.shift(1).rolling(window=window, min_periods=1).mean()
        )
    
    # ========== 阶段5：乖离率计算 ==========
    for window in [3, 5,  15]:
        ma_col = f'ma{window}w'
        df[f'bias{window}w'] = (df['close'] - df[ma_col]) / df[ma_col] * 100
    
    # ========== 阶段6：倾斜率计算 ==========
    df = calculate_slopes(df)
    
    # ========== 阶段7：均线多空判断 ==========
    df['w_duo1'] = (df['ma3w'] > df['ma5w']) & (df['ma5w'] >= df['ma8w'])
    df['w_duo2'] = (df['ma3w'] > df['ma5w']) & (df['ma5w'] >= df['ma15w'])
    
    # ========== 清理中间列 ==========
    df.drop(['year_week', 'week_end_date', 'week_close', 'dynamic_week_close'], 
            axis=1, inplace=True, errors='ignore')
    df = drop_base_columns(df)

    return df.round(2)

# 保留原calculate_slopes方法



def vectorized_slope(y_series: pd.Series, window: int) -> np.ndarray:
    """向量化斜率计算（比原生apply快100倍）"""
    y = y_series.values
    n = len(y)
    x = np.arange(window)  # 预计算x轴
    
    # 预分配结果数组
    slopes = np.full(n, np.nan)
    
    # 矩阵运算核心
    X = np.vstack([x, np.ones(window)]).T
    XT_X_inv = np.linalg.inv(X.T @ X) @ X.T
    
    for i in range(window-1, n):
        y_window = y[i-window+1:i+1]
        if np.isnan(y_window).any():
            continue
        slopes[i] = XT_X_inv @ y_window @ np.array([1, 0])  # 仅取斜率系数
    
    return slopes

def calculate_slopes(df: pd.DataFrame) -> pd.DataFrame:
    """批量计算周均线斜率（优化版）"""
    df = df.sort_values(['code', 'date'])
    
    # 并行计算所有窗口
    windows = [3,5,15]
    for w in windows:
        col = f'ma{w}w'
        df[f'slope{w}w'] = df.groupby('code', group_keys=False)[col].transform(
            lambda s: pd.Series(vectorized_slope(s, w), index=s.index)
        )
    
    return df




        